package com.highper.gateway.filter;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
import org.springframework.cloud.gateway.config.LoadBalancerProperties;
import org.springframework.cloud.gateway.filter.LoadBalancerClientFilter;
import org.springframework.cloud.netflix.ribbon.RibbonLoadBalancerClient;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.server.ServerWebExchange;

import java.net.URI;
import java.util.Objects;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR;

/**
 * @Author ZhangGJ
 * @Date 2022/02/11 11:42
 */
@Slf4j
@Configuration
public class GrayLoadBalancerClientFilter extends LoadBalancerClientFilter {

    @Autowired
    private LoadBalancerClient client;

    @Autowired
    private LoadBalancerProperties properties;

    public GrayLoadBalancerClientFilter(LoadBalancerClient loadBalancer,
            LoadBalancerProperties properties) {
        super(loadBalancer, properties);
    }

    @Override
    public ServiceInstance choose(ServerWebExchange exchange) {
        String appVersion = exchange.getRequest().getHeaders().getFirst("appVersion");
        if (StringUtils.isEmpty(appVersion)) {
            return super.choose(exchange);
        }
        if (this.loadBalancer instanceof RibbonLoadBalancerClient) {
            RibbonLoadBalancerClient client = (RibbonLoadBalancerClient) this.loadBalancer;
            String serviceId = ((URI) Objects.requireNonNull(
                    exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR))).getHost();
            return client.choose(serviceId, appVersion);
        }
        return super.choose(exchange);
    }
}
